package work.soho.proxy.utils;

import java.io.*;
import java.math.BigInteger;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;

public class RsaUtils {
    /**
     * 创建Rsa密钥对
     *
     * @return
     */
    public static KeyPair createRsa() throws NoSuchAlgorithmException {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
        keyGen.initialize(1024);
        KeyPair pair = keyGen.generateKeyPair();
        return pair;
    }

    /**
     * 创建密钥对到文件
     *
     * @param name
     * @throws NoSuchAlgorithmException
     */
    public static void randRsa2File(String name) {
        try {
            KeyPair keyPair = createRsa();
            String userHome = System.getProperty("user.home");
            String keyPath = userHome + "/.ssh";
            File file = new File(keyPath);
            if(!file.exists() || !file.isDirectory()) {
                file.mkdirs();
            }
            System.out.println(keyPath);
            //写入私钥
            File privateKey = new File(keyPath + "/" + name);
            PrintStream printStream = new PrintStream(new FileOutputStream(privateKey));
            printStream.println("-----BEGIN OPENSSH PRIVATE KEY-----");
            printStream.print(base64Encode(keyPair.getPrivate().getEncoded()));
            printStream.println("-----END OPENSSH PRIVATE KEY-----");
            printStream.close();

            //写入公钥
            File publicKey = new File(keyPath + "/" + name + ".pub");
            PrintStream publicPrintStream = new PrintStream(new FileOutputStream(publicKey));
            publicPrintStream.print("ssh-rsa " + base64Encode(keyPair.getPrivate().getEncoded()));
            publicPrintStream.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 读取指定名字公钥
     *
     * @param name
     * @return
     * @throws IOException
     */
    public static String getOrCreatePublicKey(String name) throws IOException {
        String userHome = System.getProperty("user.home");
        String keyPath = userHome + "/.ssh/" + name + ".pub";
        File file = new File(keyPath);
        if(!file.exists() || !file.isFile()) {
            randRsa2File(name);
        }
        BufferedReader reader = new BufferedReader(new FileReader(file));
        String tempString = null;
        StringBuffer stringBuffer = new StringBuffer();
        while ((tempString = reader.readLine()) != null) {
            stringBuffer.append(tempString);
        }
        reader.close();
        return stringBuffer.toString();
    }

    private static String base64Encode(byte[] bytes) {
        return Base64.getEncoder().encodeToString(bytes);
    }

    /**
     * 获取公钥签名
     *
     * @param name
     * @return
     */
    public static String md5PublicKeyByName(String name) {
        try {
            String key = getOrCreatePublicKey(name);
            MessageDigest md = MessageDigest.getInstance("MD5");
            md.update(key.getBytes());
            return new BigInteger(1, md.digest()).toString(16);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }
}
